from langchain_core.messages import HumanMessage
from langgraph.prebuilt import create_react_agent
from langchain_openai import ChatOpenAI
from langchain.tools import tool
from langmem import create_manage_memory_tool,create_search_memory_tool
from langgraph.store.memory import InMemoryStore
import os

# 设置 API 密钥
# os.environ["OPENAI_API_KEY"] = "your-api-key-here"

# 定义工具
@tool
def get_weather(city: str) -> str:
    """获取城市天气信息"""
    # 这里应该是实际的天气 API 调用
    return f"{city}的天气是晴朗，25°C"

@tool
def get_time(timezone: str) -> str:
    """获取指定时区的时间"""
    from datetime import datetime
    import pytz
    try:
        tz = pytz.timezone(timezone)
        current_time = datetime.now(tz).strftime("%Y-%m-%d %H:%M:%S")
        return f"{timezone}的当前时间是: {current_time}"
    except:
        return "时区不存在"

# 工具列表
tools = [get_weather, get_time,create_manage_memory_tool(namespace=("memories")),
         create_search_memory_tool(namespace=("memories"))]

# 创建模型（确保使用正确版本）
try:
    model = ChatOpenAI(
        model="qwen-plus",  # 或 "qwen-turbo", "qwen-max"
        api_key='sk-965dc39b016c49ecbe29de180f4db2b6',
        base_url="https://dashscope.aliyuncs.com/compatible-mode/v1",)
    # 尝试绑定工具（新版本方法）
    model_with_tools = model.bind_tools(tools)
except AttributeError:
    # 如果 bind_tools 不存在，直接使用模型
    model_with_tools = model

# 创建代理
agent = create_react_agent(model_with_tools, tools,store=InMemoryStore())

# 运行代理
config = {"configurable": {"thread_id": "test123"}}
result = agent.invoke({
    "messages": [("user", "请记住我喜欢编程")]
}, config)

print(result["messages"][-1].content)

search_result = agent.invoke({"messages":[
    HumanMessage(content="回忆一下我喜欢什么?")
]})
print(result["messages"][-1].content)